#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：auto-test 
@File    ：report_add_image.py
@Author  ：赵加响
@Date    ：2023/1/9 13:28 
@Desc    ：
"""
from time import sleep

import allure


def add_image_2_report(driver, step_name, need_sleep=True):
    """
    截图并插入测试报告
    :param driver:
    :param step_name:
    :param need_sleep:
    :return:
    """

    if need_sleep:
        sleep(2)
    attach_obj = driver.get_screenshot_as_png()
    attach_name = step_name + '.png'
    attach_type = allure.attachment_type.PNG
    allure.attach(attach_obj, attach_name, attach_type)
